# -*- coding: utf-8 -*-
import pandas as pd
import numpy as np
import matplotlib.pyplot as plt
from statsmodels.tsa.arima.model import ARIMA
from sklearn.preprocessing import MinMaxScaler
from sklearn.metrics import mean_absolute_error, mean_squared_error, r2_score
from tensorflow.keras.models import Sequential
from tensorflow.keras.layers import LSTM, Dense, Dropout
from tensorflow.keras.callbacks import EarlyStopping
import warnings

warnings.filterwarnings('ignore')


# ======================
# 1. 数据预处理模块
# ======================
def load_and_preprocess(filepath):
    """加载并预处理原始数据"""
    df = pd.read_csv(filepath)

    # 转换日期格式并排序
    df['日期'] = pd.to_datetime(df['日期'])
    df = df.sort_values('日期').reset_index(drop=True)

    # 处理交易量字符串（如"21.25M"）
    df['交易量'] = df['交易量'].str.replace('M', '').astype(float) * 1e6

    # 构造基础特征
    df['涨跌幅'] = df['涨跌幅'].str.replace('%', '').astype(float) / 100
    df['价格波动'] = (df['高'] - df['低']) / df['收盘']

    # 技术指标
    df['5日均线'] = df['收盘'].rolling(5).mean()
    df['RSI_14'] = calculate_rsi(df['收盘'], 14)

    # 滞后特征
    for lag in [1, 2, 3]:
        df[f'收盘_lag{lag}'] = df['收盘'].shift(lag)

    # 删除缺失值
    df = df.dropna()

    return df


def calculate_rsi(series, window=14):
    """计算相对强弱指数RSI"""
    delta = series.diff()
    gain = delta.where(delta > 0, 0)
    loss = -delta.where(delta < 0, 0)

    avg_gain = gain.rolling(window).mean()
    avg_loss = loss.rolling(window).mean()

    rs = avg_gain / avg_loss
    return 100 - (100 / (1 + rs))


# ======================
# 2. 特征工程模块
# ======================
def prepare_features(df, target_col='收盘'):
    """准备模型输入特征"""
    # 选择特征列
    feature_cols = ['开盘', '高', '低', '交易量', '涨跌幅', '价格波动',
                    '5日均线', 'RSI_14', '收盘_lag1', '收盘_lag2']

    # 标准化
    scaler = MinMaxScaler()
    scaled_features = scaler.fit_transform(df[feature_cols])

    # 构建时间步长数据（3天窗口）
    X, y = [], []
    for i in range(3, len(df)):
        X.append(scaled_features[i - 3:i])
        y.append(df[target_col].iloc[i])

    return np.array(X), np.array(y), scaler


# ======================
# 3. 模型构建模块
# ======================
def build_arima(train_data):
    """构建ARIMA模型"""
    model = ARIMA(train_data, order=(2, 1, 2))
    model_fit = model.fit()
    return model_fit


def build_lstm(input_shape):
    """构建LSTM模型"""
    model = Sequential([
        LSTM(128, return_sequences=True, input_shape=input_shape),
        Dropout(0.2),
        LSTM(128),
        Dense(1)
    ])
    model.compile(optimizer='adam', loss='mse')
    return model


# ======================
# 4. 训练与预测模块
# ======================
def hybrid_predict(df):
    """混合模型预测流程"""
    # 数据准备
    train_size = int(len(df) * 0.8)
    train_df, test_df = df.iloc[:train_size], df.iloc[train_size:]

    # ARIMA部分
    arima_model = build_arima(train_df['收盘'])
    arima_pred = arima_model.predict(start=len(train_df), end=len(df) - 1)

    # LSTM部分
    X_train, y_train, scaler = prepare_features(train_df)
    X_test, y_test, _ = prepare_features(test_df)

    lstm_model = build_lstm((X_train.shape[1], X_train.shape[2]))
    early_stop = EarlyStopping(monitor='val_loss', patience=5)
    lstm_model.fit(X_train, y_train,
                   epochs=50, batch_size=16,
                   validation_data=(X_test, y_test),
                   callbacks=[early_stop],
                   verbose=0)

    lstm_pred = lstm_model.predict(X_test).flatten()

    # 组合预测（权重可调）
    hybrid_pred = 0.6 * arima_pred + 0.4 * lstm_pred

    return {
        'true': y_test,
        'arima': arima_pred,
        'lstm': lstm_pred,
        'hybrid': hybrid_pred,
        'dates': test_df['日期'].iloc[3:]
    }


# ======================
# 5. 未来预测模块
# ======================
def forecast_future(df, days=30):
    """预测未来价格"""
    # 复制最后3天数据作为初始输入
    last_window = df.iloc[-3:].copy()
    future_dates = pd.date_range(df['日期'].iloc[-1] + pd.Timedelta(days=1), periods=days)

    predictions = []
    for _ in range(days):
        # 准备输入特征
        X, _, scaler = prepare_features(last_window)
        X = X[-1].reshape(1, 3, -1)  # 取最新窗口

        # LSTM预测
        lstm_pred = lstm_model.predict(X).flatten()[0]

        # ARIMA预测（需更新模型）
        arima_model = build_arima(df['收盘'])
        arima_pred = arima_model.predict(start=len(df), end=len(df))[0]

        # 组合预测
        pred_close = 0.6 * arima_pred + 0.4 * lstm_pred

        # 生成高低价（基于历史波动率）
        volatility = df['价格波动'].mean()
        pred_high = pred_close * (1 + volatility / 2)
        pred_low = pred_close * (1 - volatility / 2)

        # 更新last_window
        new_row = {
            '日期': future_dates[len(predictions)],
            '收盘': pred_close,
            '开盘': pred_close,  # 简化处理
            '高': pred_high,
            '低': pred_low,
            '交易量': df['交易量'].mean(),
            '涨跌幅': (pred_close - last_window['收盘'].iloc[-1]) / last_window['收盘'].iloc[-1]
        }
        last_window = last_window.append(new_row, ignore_index=True).iloc[-3:]

        predictions.append([pred_close, pred_high, pred_low])

    return pd.DataFrame(predictions,
                        columns=['预测收盘价', '预测最高价', '预测最低价'],
                        index=future_dates)


# ======================
# 6. 主执行流程
# ======================
if __name__ == "__main__":
    # 加载数据
    filepath = "D:/MathD/中国平安2318历史数据.csv"
    df = load_and_preprocess(filepath)

    # 训练混合模型
    results = hybrid_predict(df)

    # 评估模型
    print(f"Hybrid Model MAE: {mean_absolute_error(results['true'], results['hybrid']):.2f}")
    print(f"Hybrid Model R2: {r2_score(results['true'], results['hybrid']):.2f}")

    # 可视化历史拟合
    plt.figure(figsize=(12, 6))
    plt.plot(results['dates'], results['true'], label='Actual')
    plt.plot(results['dates'], results['hybrid'], label='Hybrid Prediction')
    plt.title('Historical Price vs Prediction')
    plt.legend()
    plt.savefig('historical_fit.png')
    plt.show()

    # 未来预测
    future_pred = forecast_future(df, days=30)
    print(future_pred.head())

    # 保存结果
    future_pred.to_csv('中国平安_20250808-20250905_预测结果.csv')

    # 可视化未来预测
    plt.figure(figsize=(12, 6))
    plt.plot(df['日期'], df['收盘'], label='Historical')
    plt.plot(future_pred.index, future_pred['预测收盘价'], label='Forecast')
    plt.fill_between(future_pred.index,
                     future_pred['预测最低价'],
                     future_pred['预测最高价'],
                     alpha=0.2)
    plt.title('30-Day Price Forecast with Confidence Interval')
    plt.legend()
    plt.savefig('future_forecast.png')
    plt.show()